"""Streamlit frontend for ChatWithRepository."""

import os
import sys
from pathlib import Path
import streamlit as st
import time
from typing import List, Dict, Any, Optional

# Add the project root to the Python path
sys.path.append(str(Path(__file__).parent))

from chat_with_repo.config import Config
from chat_with_repo.github_downloader import GitHubDownloader
from chat_with_repo.query_engine import QueryEngine
from chat_with_repo.rag_indexer import RAGIndexer

# Page configuration
st.set_page_config(
    page_title="ChatWithRepository",
    page_icon="🤖",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS
st.markdown("""
    <style>
    .main {
        padding: 0rem 1rem;
    }
    .stChat {
        height: 400px;
        overflow-y: auto;
    }
    .code-block {
        background-color: #f0f2f6;
        border-radius: 5px;
        padding: 10px;
        margin: 10px 0;
        font-family: 'Courier New', monospace;
    }
    .repository-card {
        background-color: #f8f9fa;
        border: 1px solid #dee2e6;
        border-radius: 8px;
        padding: 15px;
        margin: 10px 0;
    }
    .status-indicator {
        display: inline-block;
        width: 10px;
        height: 10px;
        border-radius: 50%;
        margin-right: 5px;
    }
    .status-active {
        background-color: #28a745;
    }
    .status-inactive {
        background-color: #dc3545;
    }
    </style>
""", unsafe_allow_html=True)

# Session state initialization
if 'messages' not in st.session_state:
    st.session_state.messages = []

if 'current_repo' not in st.session_state:
    st.session_state.current_repo = None

if 'query_engine' not in st.session_state:
    st.session_state.query_engine = None

if 'repos_status' not in st.session_state:
    st.session_state.repos_status = {}

# Initialize components
@st.cache_resource
def init_components():
    """Initialize application components."""
    try:
        Config.validate()
        downloader = GitHubDownloader()
        indexer = RAGIndexer()
        engine = QueryEngine()
        return downloader, indexer, engine
    except Exception as e:
        st.error(f"Failed to initialize components: {str(e)}")
        return None, None, None

def check_configuration():
    """Check if API keys are configured."""
    if not Config.KIMI_API_KEY:
        st.error("🔑 KIMI API Key not found!")
        st.info("Please set your KIMI_API_KEY in the .env file or run setup.")
        
        with st.expander("Setup Instructions"):
            st.code("""
            # Create .env file with:
            KIMI_API_KEY=your_kimi_api_key_here
            GITHUB_TOKEN=your_github_token_here
            """)
        return False
    return True

def get_repositories_status():
    """Get status of all repositories."""
    downloader = GitHubDownloader()
    indexer = RAGIndexer()
    
    try:
        local_repos = downloader.list_repositories()
        repos_data = []
        
        for repo_name in local_repos:
            repo_path = Path(Config.REPOS_DIR) / repo_name
            stats = indexer.get_repository_stats(str(repo_path))
            
            repos_data.append({
                'name': repo_name,
                'path': str(repo_path),
                'indexed': stats["indexed"],
                'chunks': stats["chunks"],
                'status': '✅ Indexed' if stats["indexed"] else '❌ Not Indexed'
            })
        
        return repos_data
    except Exception:
        return []

def add_repository(repo_url: str, branch: str = "main", force_reindex: bool = False):
    """Add and index a repository."""
    try:
        downloader = GitHubDownloader()
        indexer = RAGIndexer()
        
        # Check if already indexed
        if not force_reindex:
            stats = indexer.get_repository_stats(repo_url)
            if stats["indexed"]:
                st.info(f"Repository already indexed ({stats['chunks']} chunks)")
                return True
        
        # Download repository
        with st.spinner("📥 Downloading repository..."):
            repo_path = downloader.download_repository(repo_url, branch)
        
        # Index repository
        with st.spinner("🔄 Indexing repository..."):
            chunk_count = indexer.index_repository(repo_path, repo_url)
        
        st.success(f"✅ Successfully indexed {chunk_count} code chunks")
        return True
        
    except Exception as e:
        st.error(f"❌ Failed to add repository: {str(e)}")
        return False

def query_repository(query: str, repo_url: str = None) -> str:
    """Query the repository."""
    try:
        if not st.session_state.query_engine:
            st.session_state.query_engine = QueryEngine()
        
        response = st.session_state.query_engine.query_repository(query, repo_url)
        return response
        
    except Exception as e:
        return f"Error querying repository: {str(e)}"

def main():
    """Main application."""
    
    # Title and description
    st.title("🤖 ChatWithRepository")
    st.markdown("Natural language querying for GitHub repositories")
    
    # Check configuration
    if not check_configuration():
        return
    
    # Initialize components
    downloader, indexer, engine = init_components()
    
    # Sidebar
    with st.sidebar:
        st.header("📁 Repository Management")
        
        # Add repository section
        with st.expander("➕ Add Repository", expanded=True):
            repo_url = st.text_input(
                "Repository URL",
                placeholder="https://github.com/user/repository",
                help="Enter a GitHub repository URL"
            )
            
            branch = st.text_input("Branch", value="main")
            force_reindex = st.checkbox("Force reindex", help="Reindex even if already indexed")
            
            if st.button("Add Repository", type="primary"):
                if repo_url:
                    if add_repository(repo_url, branch, force_reindex):
                        st.rerun()
                else:
                    st.error("Please enter a repository URL")
        
        # Repository list
        st.header("📊 Repositories")
        repos = get_repositories_status()
        
        if repos:
            for repo in repos:
                col1, col2, col3 = st.columns([3, 1, 1])
                with col1:
                    st.write(f"**{repo['name']}**")
                with col2:
                    st.write(repo['status'])
                with col3:
                    if repo['indexed']:
                        st.write(f"{repo['chunks']} chunks")
        else:
            st.info("No repositories found. Add one above!")
    
    # Main content area
    col1, col2 = st.columns([2, 1])
    
    with col1:
        st.header("💬 Chat Interface")
        
        # Repository selection for chat
        repos = get_repositories_status()
        indexed_repos = [repo for repo in repos if repo['indexed']]
        
        if indexed_repos:
            repo_names = [repo['name'] for repo in indexed_repos]
            selected_repo = st.selectbox(
                "Select Repository",
                [None] + repo_names,
                format_func=lambda x: "All Repositories" if x is None else x
            )
            
            if selected_repo:
                for repo in indexed_repos:
                    if repo['name'] == selected_repo:
                        st.session_state.current_repo = repo['path']
                        break
        
        # Chat interface
        for message in st.session_state.messages:
            with st.chat_message(message["role"]):
                st.markdown(message["content"])
        
        # Chat input
        if prompt := st.chat_input("Ask about the repository..."):
            st.session_state.messages.append({"role": "user", "content": prompt})
            
            with st.chat_message("user"):
                st.markdown(prompt)
            
            with st.chat_message("assistant"):
                with st.spinner("🤔 Thinking..."):
                    repo_url = st.session_state.current_repo if hasattr(st.session_state, 'current_repo') else None
                    response = query_repository(prompt, repo_url)
                    st.markdown(response)
                    
                    st.session_state.messages.append({"role": "assistant", "content": response})
        
        # Clear chat button
        if st.session_state.messages:
            if st.button("🗑️ Clear Chat"):
                st.session_state.messages = []
                st.rerun()
    
    with col2:
        st.header("🎯 Quick Actions")
        
        # Quick query buttons
        st.subheader("Common Queries")
        
        quick_queries = [
            "What does the main function do?",
            "Show me error handling examples",
            "How is authentication implemented?",
            "Find database connection code",
            "Explain the API endpoints",
            "Show configuration examples",
            "Where are tests located?",
            "Find logging implementations"
        ]
        
        for query in quick_queries:
            if st.button(query, use_container_width=True):
                if st.session_state.query_engine:
                    st.session_state.messages.append({"role": "user", "content": query})
                    
                    with st.spinner("🤔 Thinking..."):
                        repo_url = st.session_state.current_repo if hasattr(st.session_state, 'current_repo') else None
                        response = query_repository(query, repo_url)
                        st.session_state.messages.append({"role": "assistant", "content": response})
                        st.rerun()
                else:
                    st.error("Query engine not initialized")
        
        # Status information
        st.subheader("📊 Status")
        
        try:
            repos = get_repositories_status()
            total_repos = len(repos)
            indexed_repos = len([r for r in repos if r['indexed']])
            total_chunks = sum([r['chunks'] for r in repos if r['indexed']])
            
            st.metric("Total Repositories", total_repos)
            st.metric("Indexed Repositories", indexed_repos)
            st.metric("Total Chunks", total_chunks)
            
        except Exception:
            st.error("Could not load status information")

if __name__ == "____main____":
    main()

# This is needed for Streamlit
if __name__ == "__main__":
    main()